import os
import sys
from numbers import Integral, Real
from collections import Iterable

import numpy as np
import numpy.random
import matplotlib
from mpl_toolkits.mplot3d import Axes3D

# Force headless backend for plotting on clusters
if "DISPLAY" not in os.environ:
    matplotlib.use('Agg')

import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cmx

import openmoc

# For Python 2.X.X
if (sys.version_info[0] == 2):
    from log import *
    from process import get_scalar_fluxes
    import checkvalue as cv
# For Python 3.X.X
else:
    from openmoc.log import *
    from openmoc.process import get_scalar_fluxes
    import openmoc.checkvalue as cv

# Store viable OpenMOC solver types for type checking
solver_types = (openmoc.Solver,)
try:
    # Try to import OpenMOC's CUDA module
    if (sys.version_info[0] == 2):
        from cuda import GPUSolver
    else:
        from openmoc.cuda import GPUSolver
    solver_types += (GPUSolver,)
except ImportError:
    pass

# Force non-interactive mode for plotting on clusters
plt.ioff()

# Default matplotlib parameters to use in all plots
matplotlib_rcparams = matplotlib.rcParamsDefault
matplotlib_rcparams['font.family'] = 'sans-serif'
matplotlib_rcparams['font.weight'] = 'normal'
matplotlib_rcparams['font.size'] = 15
matplotlib_rcparams['savefig.dpi'] = 500
matplotlib_rcparams['figure.dpi'] = 500

# A static variable for the output directory in which to save plots
subdirectory = "/plots/"

TINY_MOVE = openmoc.TINY_MOVE

if sys.version_info[0] >= 3:
    basestring = str


def plot_tracks(track_generator, get_figure=False, plot_3D=False):
    """Plot the characteristic tracks from an OpenMOC simulation.

    This method requires that tracks have been generated by a TrackGenerator.

    Parameters
    ----------
    track_generator : openmoc.TrackGenerator
        A TrackGenerator with the tracks to plot
    get_figure : bool
        Whether or not to return the Matplotlib figure

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_tracks(track_generator)

    """

    cv.check_type('track_generator', track_generator, openmoc.TrackGenerator)
    if not track_generator.containsTracks():
        py_printf('ERROR', 'Unable to plot Tracks since the track ' +
                  'generator has not yet generated tracks')

    global subdirectory, matplotlib_rcparams
    directory = openmoc.get_output_directory() + subdirectory

    # Ensure that normal settings are used even if called from ipython
    curr_rc = dict(matplotlib.rcParams)
    matplotlib.rcParams.update(curr_rc)

    # Make directory if it does not exist
    try:
        os.makedirs(directory)
    except OSError:
        pass

    py_printf('NORMAL', 'Plotting the tracks...')

    # Retrieve data from TrackGenerator
    vals_per_track = openmoc.NUM_VALUES_PER_RETRIEVED_TRACK
    num_azim = track_generator.getNumAzim()
    spacing = track_generator.getDesiredAzimSpacing()
    num_tracks = int(track_generator.getNumTracks())
    coords = track_generator.retrieveTrackCoords(num_tracks*vals_per_track)

    # Convert data to NumPy arrays
    coords = np.array(coords)
    x = coords[0::vals_per_track//2]
    y = coords[1::vals_per_track//2]
    z = coords[2::vals_per_track//2]

    # Make figure of line segments for each Track
    fig = plt.figure()
    fig.patch.set_facecolor('none')
    if plot_3D:
        ax = fig.gca(projection = '3d')
        for i in range(num_tracks):
            plt.plot(x[i*2:(i+1)*2], y[i*2:(i+1)*2], z[i*2:(i+1)*2],  'b-')
        if z.min() != z.max():
          ax.set_zlim(z.min(), z.max())
    else:
        for i in range(num_tracks):
            plt.plot(x[i*2:(i+1)*2], y[i*2:(i+1)*2], 'b-')

    plt.xlim([x.min(), x.max()])
    plt.ylim([y.min(), y.max()])


    title = 'Tracks for {0} angles and {1} cm spacing'\
        .format(num_azim, spacing)
    plt.title(title)

    # Restore settings if called from ipython
    matplotlib.rcParams.update(curr_rc)

    # Save the figure to a file or return to user
    if track_generator.getGeometry().isRootDomain():
        if get_figure:
            return fig
        else:
            filename = \
                'tracks-{1}-angles-{2}.png'.format(directory, num_azim,
                                                   spacing)
            if plot_3D:
                filename = '3d-' + filename
            fig.savefig(directory+filename, bbox_inches='tight')
            plt.close(fig)

    del coords


def plot_segments(track_generator, get_figure=False, plot_3D=False):
    """Plot the characteristic track segments from an OpenMOC simulation.

    This method requires that tracks have been generated by a TrackGenerator.
    Each segment is colored by the ID of the unique FSR it is within.

    Parameters
    ----------
    track_generator : openmoc.TrackGenerator
        A TrackGenerator with the track segments to plot
    get_figure : bool
        Whether or not to return the Matplotlib figure

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_segments(track_generator)

    """

    cv.check_type('track_generator', track_generator, openmoc.TrackGenerator)
    if not track_generator.containsTracks():
        py_printf('ERROR', 'Unable to plot Track segments since the ' +
                  'TrackGenerator has not yet generated Tracks.')

    global subdirectory, matplotlib_rcparams
    directory = openmoc.get_output_directory() + subdirectory

    # Ensure that normal settings are used even if called from ipython
    curr_rc = dict(matplotlib.rcParams)
    matplotlib.rcParams.update(curr_rc)

    # Make directory if it does not exist
    try:
        os.makedirs(directory)
    except OSError:
        pass

    py_printf('NORMAL', 'Plotting the track segments...')

    # Retrieve data from TrackGenerator
    vals_per_segment = openmoc.NUM_VALUES_PER_RETRIEVED_SEGMENT
    num_azim = track_generator.getNumAzim()
    spacing = track_generator.getDesiredAzimSpacing()
    num_segments = int(track_generator.getNumSegments())
    num_fsrs = int(track_generator.getGeometry().getNumTotalFSRs())
    coords = \
        track_generator.retrieveSegmentCoords(num_segments*vals_per_segment)

    # Convert data to NumPy arrays
    coords = np.array(coords)
    x = np.zeros(num_segments*2)
    y = np.zeros(num_segments*2)
    z = np.zeros(num_segments*2)
    fsrs = np.zeros(num_segments)

    for i in range(num_segments):
        fsrs[i] = coords[i*vals_per_segment]
        x[i*2] = coords[i*vals_per_segment+1]
        y[i*2] = coords[i*vals_per_segment+2]
        z[i*2] = coords[i*vals_per_segment+3]
        x[i*2+1] = coords[i*vals_per_segment+4]
        y[i*2+1] = coords[i*vals_per_segment+5]
        z[i*2+1] = coords[i*vals_per_segment+6]

    # Create array of equally spaced randomized floats as a color map for plots
    # Seed the NumPy random number generator to ensure reproducible color maps
    numpy.random.seed(1)
    color_map = np.linspace(0., 1., num_fsrs, endpoint=False)
    numpy.random.shuffle(color_map)

    # Make figure of line segments for each track
    fig = plt.figure()
    fig.patch.set_facecolor('none')

    # Create a color map corresponding to FSR IDs
    if plot_3D:
        ax = fig.gca(projection = '3d')
        for i in range(num_segments):
            cNorm  = colors.Normalize(vmin=0, vmax=max(color_map))
            scalarMap = cmx.ScalarMappable(norm=cNorm)
            color = scalarMap.to_rgba(color_map[int(fsrs[i]) % num_fsrs])
            plt.plot(x[i*2:(i+1)*2], y[i*2:(i+1)*2], z[i*2:(i+1)*2],  c=color)
        if z.min() != z.max():
          ax.set_zlim(z.min(), z.max())
    else:
        for i in range(num_segments):
            cNorm  = colors.Normalize(vmin=0, vmax=max(color_map))
            scalarMap = cmx.ScalarMappable(norm=cNorm)
            color = scalarMap.to_rgba(color_map[int(fsrs[i]) % num_fsrs])
            plt.plot(x[i*2:(i+1)*2], y[i*2:(i+1)*2], c=color)

    plt.xlim([x.min(), x.max()])
    plt.ylim([y.min(), y.max()])

    suptitle = 'Segments ({0} angles, {1} cm spacing)'.format(num_azim,
                                                              spacing)
    title = 'z = {0}'.format(z[0])
    plt.suptitle(suptitle)
    plt.title(title)

    # Restore settings if called from ipython
    matplotlib.rcParams.update(curr_rc)

    if track_generator.getGeometry().isRootDomain():
        if get_figure:
            return fig
        else:
            filename = 'segments-{0}-angles-{1}-spacing'.format(num_azim,
                                                                spacing)
            filename = '{0}-z-{1}.png'.format(filename, z[0])
            if plot_3D:
                filename = '3d-' + filename
            fig.savefig(directory+filename, bbox_inches='tight')
            plt.close(fig)


def plot_materials(geometry, gridsize=250, xlim=None, ylim=None, zlim=None,
                   plane='xy', offset=None, get_figure=False,
                   library='matplotlib'):
    """Plot a color-coded 2D surface plot of the materials in the geometry.

    The geometry must be initialized with materials, cells, universes and
    lattices before being passed into this routine.

    Parameters
    ----------
    geometry : openmoc.Geometry
        An OpenMOC geometry
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_materials(geometry)

    """

    cv.check_type('geometry', geometry, openmoc.Geometry)

    py_printf('NORMAL', 'Plotting the materials...')

    # Create an array of random integer colors for each Material
    materials = geometry.getAllMaterials()
    num_materials = len(materials)
    colors = np.arange(num_materials)
    numpy.random.seed(1)
    numpy.random.shuffle(colors)
    for i, material_id in enumerate(materials):
        materials[material_id] = colors[i]

    # Initialize plotting parameters
    plot_params = PlotParams()
    plot_params.geometry = geometry
    plot_params.domain_type = 'material'
    plot_params.gridsize = gridsize
    plot_params.library = library
    plot_params.xlim = xlim
    plot_params.ylim = ylim
    plot_params.zlim = zlim
    plot_params.plane = plane
    plot_params.offset = offset
    plot_params.suptitle = 'Materials'
    if plane == 'xy':
        plot_params.title = 'z = {0}'.format(plot_params.offset)
        plot_params.filename = 'materials-z-{0}'.format(plot_params.offset)
    elif plane == 'xz':
        plot_params.title = 'y = {0}'.format(plot_params.offset)
        plot_params.filename = 'materials-y-{0}'.format(plot_params.offset)
    elif plane == 'yz':
        plot_params.title = 'x = {0}'.format(plot_params.offset)
        plot_params.filename = 'materials-x-{0}'.format(plot_params.offset)
    plot_params.interpolation = 'nearest'
    plot_params.vmin = 0
    plot_params.vmax = num_materials

    # Plot a 2D color map of the Materials
    figures = plot_spatial_data(materials, plot_params, get_figure)

    # Return the figure to the user if requested
    if plot_params.geometry.isRootDomain():
        if get_figure:
            return figures[0]


def plot_cells(geometry, gridsize=250, xlim=None, ylim=None, zlim = None,
               plane='xy', offset=None, get_figure=False,
               library='matplotlib'):
    """Plots a color-coded 2D surface plot of the cells in the geometry.

    The geometry must be initialized with materials, cells, universes and
    lattices before being passed into this routine.

    Parameters
    ----------
    geometry : openmoc.Geometry
        An OpenMOC geometry
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_cells(geometry)

    """

    cv.check_type('geometry', geometry, openmoc.Geometry)

    py_printf('NORMAL', 'Plotting the cells...')

    # Create an array of random integer colors for each Cell
    cells = geometry.getAllMaterialCells()
    num_cells = len(cells)
    colors = np.arange(num_cells)
    numpy.random.seed(1)
    numpy.random.shuffle(colors)
    for i, cell_id in enumerate(cells):
        cells[cell_id] = colors[i]

    # Initialize plotting parameters
    plot_params = PlotParams()
    plot_params.geometry = geometry
    plot_params.domain_type = 'cell'
    plot_params.gridsize = gridsize
    plot_params.library = library
    plot_params.xlim = xlim
    plot_params.ylim = ylim
    plot_params.zlim = zlim
    plot_params.plane = plane
    plot_params.offset = offset
    plot_params.suptitle = 'Cells'
    if plane == 'xy':
        plot_params.title = 'z = {0}'.format(plot_params.offset)
        plot_params.filename = 'cells-z-{0}'.format(plot_params.offset)
    elif plane == 'xz':
        plot_params.title = 'y = {0}'.format(plot_params.offset)
        plot_params.filename = 'cells-y-{0}'.format(plot_params.offset)
    elif plane == 'yz':
        plot_params.title = 'x = {0}'.format(plot_params.offset)
        plot_params.filename = 'cells-x-{0}'.format(plot_params.offset)
    plot_params.interpolation = 'nearest'
    plot_params.vmin = 0
    plot_params.vmax = num_cells

    # Plot a 2D color map of the Cells
    figures = plot_spatial_data(cells, plot_params, get_figure)

    # Return the figure to the user if requested
    if plot_params.geometry.isRootDomain():
        if get_figure:
            return figures[0]


def plot_flat_source_regions(geometry, gridsize=250, xlim=None, ylim=None,
                             zlim=None, plane='xy', offset=0, centroids=False,
                             marker_type='o', marker_size=2, get_figure=False,
                             library='matplotlib'):
    """Plots a color-coded 2D surface plot of the FSRs in the geometry.

    The geometry must be initialized with materials, cells, universes and
    lattices before being passed into this routine.

    Parameters
    ----------
    geometry : openmoc.Geometry
        An OpenMOC geometry
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    centroids : bool
        Whether to plot the FSR centroids (False by default)
    marker_type : str
        The marker type to use for FSR centroids ('o' by default)
    marker_size : Integral
        The marker size to use for FSR centroids (2 by default)
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_flat_source_regions(geometry)

    """

    cv.check_type('geometry', geometry, openmoc.Geometry)
    cv.check_type('centroids', centroids, bool)
    cv.check_type('marker_type', marker_type, basestring)
    cv.check_value('marker_type', marker_type,
                   tuple(matplotlib.markers.MarkerStyle().markers.keys()))
    cv.check_type('marker_size', marker_size, Real)
    cv.check_greater_than('marker_size', marker_size, 0)

    if geometry.getNumTotalFSRs() == 0:
        py_printf('ERROR', 'Unable to plot the flat source regions ' +
                  'since no tracks have been generated.')

    py_printf('NORMAL', 'Plotting the flat source regions...')

    global subdirectory, matplotlib_rcparams
    directory = openmoc.get_output_directory() + subdirectory

    num_fsrs = int(geometry.getNumTotalFSRs())
    fsrs_to_fsrs = np.arange(num_fsrs, dtype=np.int64)
    fsrs_to_fsrs = _colorize(fsrs_to_fsrs, num_fsrs)

    # Initialize plotting parameters
    plot_params = PlotParams()
    plot_params.geometry = geometry
    plot_params.gridsize = gridsize
    plot_params.library = library
    plot_params.xlim = xlim
    plot_params.ylim = ylim
    plot_params.zlim = zlim
    plot_params.plane = plane
    plot_params.offset = offset
    plot_params.suptitle = 'Flat Source Regions'
    if plane == 'xy':
        plot_params.title = 'z = {0}'.format(plot_params.offset)
        plot_params.filename = 'flat-source-regions-z-{0}'\
            .format(plot_params.offset)
    elif plane == 'xz':
        plot_params.title = 'y = {0}'.format(plot_params.offset)
        plot_params.filename = 'flat-source-regions-y-{0}'\
            .format(plot_params.offset)
    elif plane == 'yz':
        plot_params.title = 'x = {0}'.format(plot_params.offset)
        plot_params.filename = 'flat-source-regions-x-{0}'\
            .format(plot_params.offset)
    plot_params.interpolation = 'nearest'
    plot_params.vmin = 0
    plot_params.vmax = num_fsrs

    # Plot a 2D color map of the flat source regions
    figures = plot_spatial_data(fsrs_to_fsrs, plot_params, get_figure=True)

    if plot_params.geometry.isRootDomain():

        fig = figures[0]

        # Plot centroids on top of 2D flat source region color map
        if centroids:

            # Populate a NumPy array with the FSR centroid coordinates
            centroids = np.zeros((num_fsrs, 2), dtype=np.float)
            for fsr_id in range(num_fsrs):
                coords = geometry.getGlobalFSRCentroidData(fsr_id)
                if plane == 'xy':
                    centroids[fsr_id,:] = [coords[0], coords[1]]
                elif plane == 'xz':
                    centroids[fsr_id,:] = [coords[0], coords[2]]
                elif plane == 'yz':
                    centroids[fsr_id,:] = [coords[1], coords[2]]

            # Plot centroids on figure using matplotlib
            if library == 'pil':

                # Retrieve the plot bounds
                coords = _get_pixel_coords(plot_params)
                r = marker_size

                # Open a PIL ImageDraw portal on the Image object
                from PIL import ImageDraw
                draw = ImageDraw.Draw(fig)

                for fsr_id in range(num_fsrs):
                    # Retrieve the pixel coordinates for this centroid
                    coord1, coord2 = centroids[fsr_id,:]

                    # Only plot centroid if it is within the plot bounds
                    if coord1 < coords['bounds'][0] or \
                        coord1 > coords['bounds'][1]:
                        continue
                    elif coord2 < coords['bounds'][2] or \
                        coord2 > coords['bounds'][3]:
                        continue

                    # Transform the centroid into pixel coordinates
                    x = int((coord1-coords['dim1'][1]) / \
                        (coords['dim1'][1]-coords['dim1'][0]))
                    y = int((coord2-coords['dim2'][1]) / \
                        (coords['dim2'][1]-coords['dim2'][0]))

                    # Draw circle for this centroid on the image
                    draw.ellipse((x-r, y-r, x+r, y+r), fill=(0, 0, 0))

            # Plot centroids on figure using PIL
            else:
                plt.scatter(centroids[:,0], centroids[:,1], color='k',
                            marker=marker_type, s=marker_size)

        # Return the figure to the user if requested
        if get_figure:
            return figures[0]
        # Set the plot title and save the figure
        else:
            plot_filename = directory + plot_params.filename + \
                plot_params.extension

            if library == 'pil':
                fig.save(plot_filename)
            else:
                fig.savefig(plot_filename, bbox_inches='tight')
                plt.close(fig)


def plot_cmfd_cells(geometry, cmfd, gridsize=250, xlim=None, ylim=None,
                    zlim=None, plane='xy', offset=0, get_figure=False,
                    library='matplotlib'):
    """Plots a color-coded 2D surface plot of the CMFD cells in a geometry.

    The geometry must be initialized with materials, cells, universes and
    lattices before being passed into this method. Plotting the CMFD cells
    requires that segments must have been created for the geometry and FSR IDs
    assigned to regions.

    Parameters
    ----------
    geometry : openmoc.Geometry
        An OpenMOC geometry
    cmfd : openmoc.Cmfd
        An OpenMOC cmfd
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_cmfd_cells(geometry, cmfd)

    """

    cv.check_type('geometry', geometry, openmoc.Geometry)
    cv.check_type('cmfd', cmfd, openmoc.Cmfd)

    py_printf('NORMAL', 'Plotting the CMFD cells...')

    # Create a NumPy array to map FSRs to CMFD cells
    num_fsrs = int(geometry.getNumTotalFSRs())
    fsrs_to_cmfd_cells = np.zeros(num_fsrs, dtype=np.int64)
    for fsr_id in range(num_fsrs):
        fsrs_to_cmfd_cells[fsr_id] = cmfd.convertGlobalFSRIdToCmfdCell(fsr_id)

    # Assign random color scheme to CMFD cells
    num_cmfd_cells = cmfd.getNumCells()
    fsrs_to_cmfd_cells = _colorize(fsrs_to_cmfd_cells, num_cmfd_cells)

    # Initialize plotting parameters
    plot_params = PlotParams()
    plot_params.geometry = geometry
    plot_params.gridsize = gridsize
    plot_params.library = library
    plot_params.xlim = xlim
    plot_params.ylim = ylim
    plot_params.zlim = zlim
    plot_params.plane = plane
    plot_params.offset = offset
    plot_params.suptitle = 'CMFD Cells'
    if plane == 'xy':
        plot_params.title = 'z = {0}'.format(plot_params.offset)
        plot_params.filename = 'cmfd-cells-z-{0}'.format(plot_params.offset)
    elif plane == 'xz':
        plot_params.title = 'y = {0}'.format(plot_params.offset)
        plot_params.filename = 'cmfd-cells-y-{0}'.format(plot_params.offset)
    elif plane == 'yz':
        plot_params.title = 'x = {0}'.format(plot_params.offset)
        plot_params.filename = 'cmfd-cells-x-{0}'.format(plot_params.offset)
    plot_params.interpolation = 'nearest'
    plot_params.vmin = 0
    plot_params.vmax = num_cmfd_cells

    # Plot the CMFD cells
    figures = plot_spatial_data(fsrs_to_cmfd_cells, plot_params, get_figure)

    # Return the figure to the user if requested
    if plot_params.geometry.isRootDomain():
        if get_figure:
            return figures[0]


def plot_spatial_fluxes(solver, energy_groups=[1], norm=False, gridsize=250,
                        xlim=None, ylim=None, zlim=None, plane='xy', offset=0,
                        get_figure=False, library='matplotlib'):
    """Plot a color-coded 2D surface plot of the FSR scalar fluxes for one or
    more energy groups.

    The solver must have converged the FSR sources before calling this routine.

    Parameters
    ----------
    solver : openmoc.Solver
        An OpenMOC solver used to compute the flux
    energy_groups : Iterable of Integral
        The energy groups to plot (starting at 1 for the highest energy)
    norm : bool, optional
        Whether to normalize the fluxes to the mean (False by default)
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : list of matplotlib.Figure or None
        The Matplotlib figures are returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_spatial_fluxes(solver, energy_groups=[1,7])

    """

    global solver_types
    cv.check_type('solver', solver, solver_types)
    cv.check_type('energy_groups', energy_groups, Iterable, Integral)

    py_printf('NORMAL', 'Plotting the FSR scalar fluxes...')

    # Initialize plotting parameters
    geometry = solver.getGeometry()
    plot_params = PlotParams()
    plot_params.geometry = geometry
    plot_params.offset = offset
    plot_params.gridsize = gridsize
    plot_params.library = library
    plot_params.xlim = xlim
    plot_params.ylim = ylim
    plot_params.zlim = zlim
    plot_params.plane = plane
    plot_params.offset = offset
    plot_params.colorbar = True
    plot_params.cmap = plt.get_cmap('jet')
    plot_params.norm = norm

    # Get array of FSR energy-dependent fluxes
    fluxes = get_scalar_fluxes(solver)

    # Initialize an empty list of Matplotlib figures if requestd by the user
    figures = []

    # Loop over all energy group and create a plot
    for index, group in enumerate(energy_groups):
        plot_params.suptitle = 'FSR Scalar Flux (Group {0})'.format(group)
        if plane == 'xy':
            plot_params.title = 'z = {0}'.format(offset)
            plot_params.filename = 'fsr-flux-group-{0}-z-{1}'.format(group,
                                                                     offset)
        elif plane == 'xz':
            plot_params.title = 'y = {0}'.format(offset)
            plot_params.filename = 'fsr-flux-group-{0}-y-{1}'.format(group,
                                                                     offset)
        elif plane == 'yz':
            plot_params.title = 'x = {0}'.format(offset)
            plot_params.filename = 'fsr-flux-group-{0}-x-{1}'.format(group,
                                                                     offset)

        fig = plot_spatial_data(fluxes[:,index], plot_params, get_figure)

        if plot_params.geometry.isRootDomain():
            if get_figure:
                figures.append(fig[0])

    # Return figures if requested by the user
    if get_figure:
        return figures


def plot_energy_fluxes(solver, fsrs, group_bounds=None, norm=True,
                       loglog=True, get_figure=False):
    """Plot the scalar flux vs. energy for one or more FSRs.

    The Solver must have converged the FSR sources before calling this routine.
    The routine will generate a step plot of the flux across each energy group.

    An optional parameter for the energy group bounds may be input. The group
    bounds should be input in increasing order of energy. If group bounds are
    not specified, the routine will use equal width steps for each energy group.

    Parameters
    ----------
    solver : openmoc.Solver
        An OpenMOC solver used to compute the flux
    fsrs : Iterable of Integral
        The FSRs for which to plot the flux
    group_bounds : Iterable of Real or None, optional
        The bounds of the energy groups
    norm : bool, optional
        Whether to normalize the fluxes (True by default)
    loglog : bool
        Whether to use a log scale on the x- and y-axes (True by default)
    get_figure : bool
        Whether to return the Matplotlib figure

    Returns
    -------
    fig : list of matplotlib.Figure or None
        The Matplotlib figures are returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_energy_fluxes(solver, fsrs=[1,5,20], \
                                               group_bounds=[0., 0.625, 2e7])

    """

    global solver_types
    cv.check_type('solver', solver, solver_types)
    cv.check_type('fsrs', fsrs, Iterable, Integral)
    cv.check_type('norm', norm, bool)
    cv.check_type('loglog', loglog, bool)
    cv.check_type('get_figure', get_figure, bool)

    geometry = solver.getGeometry()
    num_groups = geometry.getNumEnergyGroups()

    if group_bounds:
        cv.check_type('group_bounds', group_bounds, Iterable, Real)
        if not all(low < up for low, up in zip(group_bounds, group_bounds[1:])):
            py_printf('ERROR', 'Unable to plot the flux vs. energy since ' +
                      'the group bounds are not monotonically increasing')
    else:
        group_bounds = np.arange(num_groups+1, dtype=np.int)
        loglog = False

    py_printf('NORMAL', 'Plotting the scalar fluxes vs. energy...')

    global subdirectory, matplotlib_rcparams
    directory = openmoc.get_output_directory() + subdirectory

    # Ensure that normal settings are used even if called from ipython
    curr_rc = dict(matplotlib.rcParams)
    matplotlib.rcParams.update(curr_rc)

    # Make directory if it does not exist
    try:
        os.makedirs(directory)
    except OSError:
        pass

    # Compute difference in energy bounds for each group
    group_deltas = np.ediff1d(group_bounds)
    group_bounds = np.flipud(group_bounds)
    group_deltas = np.flipud(group_deltas)

    # Initialize an empty list of Matplotlib figures if requestd by the user
    figures = []

    # Iterate over all flat source regions
    for fsr in fsrs:

        # Allocate memory for an array of this FSR's fluxes
        fluxes = np.zeros(num_groups, dtype=np.float)

        # Extract the flux in each energy group
        for group in range(num_groups):
            fluxes[group] = solver.getFlux(fsr, group+1)

        # Normalize fluxes to the total integrated flux
        if norm:
            fluxes /= np.sum(group_deltas * fluxes)

        # Initialize a separate plot for this FSR's fluxes
        fig = plt.figure()
        fig.patch.set_facecolor('none')

        # Draw horizontal/vertical lines on the plot for each energy group
        for group in range(num_groups):

            # Horizontal line
            if loglog:
                plt.loglog(group_bounds[group:group+2], [fluxes[group]]*2,
                           linewidth=3, c='b', label='openmoc', linestyle='-')
            else:
                plt.plot(group_bounds[group:group+2], [fluxes[group]]*2,
                         linewidth=3, c='b', label='openmoc', linestyle='-')

            # Vertical lines
            if group < num_groups - 1:
                if loglog:
                    plt.loglog([group_bounds[group+1]]*2, fluxes[group:group+2],
                               c='b', linestyle='--')
                else:
                    plt.plot([group_bounds[group+1]]*2, fluxes[group:group+2],
                             c='b', linestyle='--')

        plt.xlabel('Energy')
        plt.ylabel('Flux')
        plt.xlim((min(group_bounds), max(group_bounds)))
        plt.grid()
        plt.title('FSR {0} Flux ({1} groups)'.format(fsr, num_groups))

        # Save the figure to a file or return to user if requested
        if geometry.isRootDomain():
            if get_figure:
                figures.append(fig)
            else:
                filename = 'flux-fsr-{0}.png'.format(fsr)
                plt.savefig(directory+filename, bbox_inches='tight')
                plt.close(fig)

    # Restore settings if called from ipython
    matplotlib.rcParams.update(curr_rc)

    # Return the figures if requested by user
    if get_figure:
        return figures


def plot_fission_rates(solver, nu=False, norm=False, transparent_zeros=True,
                       gridsize=250, xlim=None, ylim=None, zlim=None,
                       plane='xy', offset=0, get_figure=False,
                       library='matplotlib'):
    """Plot a color-coded 2D surface plot representing the FSR fission rates.

    The Solver must have converged the FSR sources before calling this routine.
    The fission rates are computed as the volume-averaged and energy-integrated
    fission rates in each FSR.

    Parameters
    ----------
    solver : openmoc.Solver
        An OpenMOC solver used to compute the flux
    nu : bool
        Whether use the nu-fission rates instead of the fission rates
        (False by default)
    norm : bool
        Whether to normalize the fission rates (False by default)
    transparent_zeros : bool
        Whether to make all non-fissionable FSRs transparent (True by default)
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

         >>> openmoc.plotter.plot_fission_rates(solver)

    """

    global solver_types
    cv.check_type('solver', solver, solver_types)

    py_printf('NORMAL', 'Plotting the flat source region fission rates...')

    # Compute the volume-weighted fission rates for each FSR
    geometry = solver.getGeometry()
    fission_rates = solver.computeFSRFissionRates(
        int(geometry.getNumTotalFSRs()), nu)

    # Initialize plotting parameters
    plot_params = PlotParams()
    plot_params.geometry = geometry
    plot_params.gridsize = gridsize
    plot_params.library = library
    plot_params.xlim = xlim
    plot_params.ylim = ylim
    plot_params.zlim = zlim
    plot_params.plane = plane
    plot_params.offset = offset
    plot_params.suptitle = 'Flat Source Region Fission Rates'
    if plane == 'xy':
        plot_params.title = 'z = {0}'.format(offset)
        plot_params.filename = 'fission-rates-z-{0}.png'.format(offset)
    elif plane == 'xz':
        plot_params.title = 'y = {0}'.format(offset)
        plot_params.filename = 'fission-rates-y-{0}.png'.format(offset)
    elif plane == 'yz':
        plot_params.title = 'x = {0}'.format(offset)
        plot_params.filename = 'fission-rates-x-{0}.png'.format(offset)
    plot_params.transparent_zeros = True
    plot_params.colorbar = True
    plot_params.cmap = plt.get_cmap('jet')
    plot_params.norm = norm

    # Plot the fission rates
    figures = plot_spatial_data(fission_rates, plot_params, get_figure)

    # Return the figure if requested by user
    if plot_params.geometry.isRootDomain():
        if get_figure:
            return figures[0]


def plot_eigenmode_fluxes(iramsolver, eigenmodes=[], energy_groups=[1],
                          norm=False, gridsize=250, xlim=None, ylim=None,
                          zlim=None, plane='xy', offset=0, get_figure=False,
                          library='matplotlib'):

    """Plot the color-coded 2D surface plot of FSR scalar fluxes for one or
    more eigenmodes from an IRAMSolver.

    The IRAMSolver must have computed the eigenmodes before calling this routine.

    Parameters
    ----------
    iramsolver : openmoc.krylov.IRAMSolver
        An OpenMOC IRAM solver used to compute the flux eigenmodes
    eigenmodes : Iterable of Integral
        The indexes of the eigenmodes to plot, 1 is the first mode
    energy_groups : Iterable of Integral
        The energy groups to plot (starting at 1 for the highest energy)
    norm : bool
        Whether to normalize the fission rates (False by default)
    gridsize : Integral, optional
        The number of grid cells for the plot (250 by default)
    xlim : 2-tuple of Real, optional
        The minimim/maximum x-coordinates
    ylim : 2-tuple of Real, optional
        The minimim/maximum y-coordinates
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    get_figure : bool
        Whether to return the Matplotlib figure (only if library='matplotlib')
    library : {'matplotlib', 'pil'}
        The plotting library to use

    Returns
    -------
    fig : list of matplotlib.Figure or None
        The Matplotlib figures are returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_eigenmode_fluxes(iramsolver, energy_groups=[3])

    """

    cv.check_type('iramsolver', iramsolver, openmoc.krylov.IRAMSolver)
    cv.check_type('eigenmodes', eigenmodes, Iterable, Integral)

    py_printf('NORMAL', 'Plotting the eigenmode fluxes...')

    global subdirectory
    directory = openmoc.get_output_directory() + subdirectory

    # Make directory if it does not exist
    try:
        os.makedirs(directory)
    except OSError:
        pass

    # Extract the MOC Solver from the IRAMSolver
    moc_solver = iramsolver._moc_solver

    # Initialize a list of figures to return to user if requested
    figures = []

    # Loop over each eigenmode
    for mode in eigenmodes:

        # Extract the eigenvector for this eigenmode from the IRAMSolver
        eigenvec = iramsolver._eigenvectors[:,mode-1]

        # Convert it into a form that SWIG will be happy with
        eigenvec = np.squeeze(np.ascontiguousarray(eigenvec))
        eigenvec = np.real(eigenvec).astype(iramsolver._precision)

        # Ensure the primary eigenvector is positive
        if mode == 1:
            eigenvec = np.abs(eigenvec)

        # Insert eigenvector into MOC Solver object
        moc_solver.setFluxes(eigenvec)

        # Set subdirectory folder for this eigenmode
        num_digits = len(str(max(eigenmodes)))
        subdirectory = \
            '/plots/eig-{0}-flux/'.format(str(mode).zfill(num_digits))

        # Plot this eigenmode's spatial fluxes
        fig = plot_spatial_fluxes(moc_solver, energy_groups, norm, gridsize,
                                  xlim, ylim, zlim, plane, offset, get_figure,
                                  library)

        if get_figure:
            figures.append(fig[0])

    # Reset global subdirectory
    subdirectory = '/plots/'

    # Return Matplotlib figures if requested by user
    if get_figure:
        return figures


def plot_spatial_data(domains_to_data, plot_params, get_figure=False):
    """Plots a color-coded 2D surface plot of arbitrary data mapped to each
    spatial domain in the geometry.

    The plot_params must include the domain type encoded in domains_to_data.

    Parameters
    ----------
    domains_to_data : dict or numpy.ndarray or pandas.DataFrame
        A mapping between spatial domain IDs and numerical data to plot.
        In the case of 'material' and 'cell' domain types, domains_to_data must
        be a Python dictionary with keys representing material/cell IDs and
        values of the data to plot. In the case of the 'fsr' domain type,
        domains_to_data may be a dictionary, NumPy array or a Pandas DataFrame
        indexed by FSR ID and values of the data to plot.
    plot_params : openmoc.plotter.PlotParams
        The plotting parameters
    get_figure : bool, optional
        Whether to return the Matplotlib figures (False by default)

    Returns
    -------
    fig : list of matplotlib.Figure or None
        The Matplotlib figures are returned if get_figure is True

    Examples
    --------
    A user may invoke this function from an Python file as follows:
        >>> fsrs_to_data = numpy.random.rand(geometry.getNumFSRS())
        >>> plot_params = PlotParams()
        >>> plot_params.domain_type = 'fsr'
        >>> openmoc.plotter.plot_spatial_data(fsrs_to_data, plot_params)

    """

    cv.check_type('plot_params', plot_params, PlotParams)

    # Determine the number of domains
    if plot_params.domain_type == 'material':
        num_domains = len(plot_params.geometry.getAllMaterials())
    elif plot_params.domain_type == 'cell':
        num_domains = len(plot_params.geometry.getAllMaterialCells())
    else:
        num_domains = int(plot_params.geometry.getNumTotalFSRs())

    if isinstance(domains_to_data, (np.ndarray, dict)):
        pandas_df = False
        if len(domains_to_data) != num_domains:
            py_printf('ERROR', 'The domains_to_data array is length %d but ' +
                      'there are %d domains', len(domains_to_data), num_domains)
    elif 'DataFrame' in str(type(domains_to_data)):
        pandas_df = True
        if len(domains_to_data) != plot_params.geometry.getNumTotalFSRs():
            py_printf('ERROR', 'The domains_to_data DataFrame is length %d ' +
                      'but there are %d domains in the Geometry',
                      len(domains_to_data), num_domains)
    else:
        py_printf('ERROR', 'Unable to plot spatial data since ' +
                  'domains_to_data is not a dict, array or DataFrame')

    global subdirectory, matplotlib_rcparams
    directory = openmoc.get_output_directory() + subdirectory

    # Make directory if it does not exist
    try:
        os.makedirs(directory)
    except OSError:
        pass

    # Retrieve the pixel coordinates
    coords = _get_pixel_coords(plot_params)

    # Query the geometry for the data on the spatial grid
    domains = plot_params.geometry.getSpatialDataOnGrid(
        coords['dim1'], coords['dim2'], offset=plot_params.offset,
        plane=plot_params.plane, domain_type=plot_params.domain_type)
    domains = np.reshape(domains, tuple(([plot_params.gridsize]*2)))
    domains[domains == np.nan] = -1

    # Make domains-to-data array 2D to mirror a Pandas DataFrame
    if isinstance(domains_to_data, np.ndarray):
        domains_to_data.shape += (1,)

    # Determine the number of plots to generate
    if pandas_df or isinstance(domains_to_data, np.ndarray):
        num_plots = domains_to_data.shape[1]
    else:
        num_plots = int(len(domains_to_data) / num_domains)

    # Initialize a list of Matplotlib figures to return to user if requested
    figures = []

    # Loop over all columns in NumPy array or Pandas DataFrame input
    for i in range(num_plots):

        # Use domain IDs to appropriately index into FSR data
        # If domains-to-data was input as a Pandas DataFrame
        if pandas_df:
            surface = domains_to_data.ix[:,i].values
            surface = surface.take(domains.flatten())
        # If domains-to-data was input as a NumPy array
        elif isinstance(domains_to_data, np.ndarray):
            surface = domains_to_data.take(domains.flatten())
        # If domains-to-data was input as a Python dictionary
        else:
            surface = np.zeros(domains.shape, dtype=np.float)
            for domain_id in domains_to_data:
                indices = np.where(domains == domain_id)
                surface[indices] = domains_to_data[domain_id]

        # Reshape data to 2D array for Matplotlib image plot
        surface.shape = (plot_params.gridsize, plot_params.gridsize)

        # Normalize data to maximum if requested
        if plot_params.norm:
            surface /= np.max(surface)

        # Set zero data entries to NaN so Matplotlib will make them transparent
        if plot_params.transparent_zeros:
            indices = np.where(surface == 0.0)
            surface[indices] = np.nan

        # Color "bad" numbers (ie, NaN, INF) with transparent pixels
        if plot_params.cmap:
            plot_params.cmap.set_bad(alpha=0.0)

        # Create plot filename
        plot_filename = directory + plot_params.filename

        # If input was Pandas DataFrame, append column name to filename
        if pandas_df:
            plot_filename += '-{0}'.format(domains_to_data.columns[i])

        # Append file extension (e.g., '.png', '.ppm') to filename
        plot_filename += plot_params.extension

        # Use Python Imaging Library (PIL) to plot 2D color map of domain data
        if plot_params.library == 'pil':
            img = _get_pil_image(np.flipud(surface), plot_params)

            if get_figure:
                figures.append(img)
            else:
                img.save(plot_filename)

        # Use Matplotlib to plot 2D color map of domain data
        else:

            # Ensure that normal settings are used even if called from ipython
            curr_rc = dict(matplotlib.rcParams)
            matplotlib.rcParams.update(curr_rc)

            fig = plt.figure()
            fig.patch.set_facecolor('none')
            plt.imshow(np.flipud(surface), extent=coords['bounds'],
                       interpolation=plot_params.interpolation,
                       vmin=plot_params.vmin, vmax=plot_params.vmax,
                       cmap=plot_params.cmap)

            if plot_params.colorbar:
                plt.colorbar()
            if plot_params.title:
                plt.title(plot_params.title)

            if plot_params.suptitle:
                # If input was Pandas DataFrame, append column name to suptitle
                if pandas_df:
                    suptitle = plot_params.suptitle
                    suptitle += ' ({0})'.format(domains_to_data.columns[i])
                else:
                    suptitle = plot_params.suptitle

                plt.suptitle(suptitle)

            # If the user requested the Matplotlib figure handles for further
            # specialization, append the handle to this figure to a list
            if plot_params.geometry.isRootDomain():
                if get_figure:
                    figures.append(fig)

                # Otherwise, save this Matplotlib figure
                else:
                    fig.savefig(plot_filename, bbox_inches='tight')
                    plt.close()

            # Restore settings if called from ipython
            matplotlib.rcParams.update(curr_rc)

    # Return Matplotlib figures if requested by user
    if get_figure:
        return figures

    # Delete the memory allocated for the domains std::vector
    del domains


def plot_quadrature(solver, get_figure=False):
    """Plots the quadrature set used for an OpenMOC simulation.

    This method requires that tracks have been generated by a TrackGenerator
    owned by the solver.

    Paramters
    ---------
    solver : openmoc.Solver
        An OpenMOC solver with a trackgenerator
    get_figure : bool
        Whether to return the Matplotlib figure

    Returns
    -------
    fig : matplotlib.Figure or None
        The Matplotlib figure is returned if get_figure is True

    Examples
    --------
    A user may invoke this routine from an OpenMOC Python file as follows:

        >>> openmoc.plotter.plot_quadrature(solver)

    """

    global solver_types
    cv.check_type('solver', solver, solver_types)

    py_printf('NORMAL', 'Plotting the quadrature...')

    global subdirectory, matplotlib_rcparams
    directory = openmoc.get_output_directory() + subdirectory

    # Ensure that normal settings are used even if called from ipython
    curr_rc = dict(matplotlib.rcParams)
    matplotlib.rcParams.update(curr_rc)

    # Make directory if it does not exist
    try:
        os.makedirs(directory)
    except OSError:
        pass

    # Retrieve data from TrackGenerator
    track_generator = solver.getTrackGenerator()
    quad = track_generator.getQuadrature()
    num_azim = track_generator.getNumAzim()
    azim_spacing = track_generator.getDesiredAzimSpacing()
    num_polar_2 = int(quad.getNumPolarAngles() / 2)
    phis = np.zeros(num_azim//4)
    thetas = np.zeros(num_polar_2)

    # Get the polar angles
    for p in range(num_polar_2):
        thetas[p] = np.arcsin(quad.getSinTheta(0,p))

    # Get the azimuthal angles
    for a in range(int(num_azim / 4)):
        phis[a] = quad.getPhi(a)

    # Make a 3D figure
    fig = plt.figure()
    fig.patch.set_facecolor('none')
    ax = fig.gca(projection ='3d')

    # Plot a wire mesh on one octant of the unit sphere
    u = np.linspace(0, np.pi/2.0, 100)
    v = np.linspace(0, np.pi/2.0, 100)
    x = np.outer(np.cos(u), np.sin(v))
    y = np.outer(np.sin(u), np.sin(v))
    z = np.outer(np.ones(np.size(u)), np.cos(v))
    ax.plot_wireframe(x, y, z, rstride=5, cstride=5, color='k', linewidth=0.1)

    # Plot the quadrature points on the octant unit sphere
    for a in range(int(num_azim / 4)):
        for p in range(num_polar_2):
            ax.scatter(np.cos(phis[a]) * np.sin(thetas[p]), np.sin(phis[a]) *
                       np.sin(thetas[p]), np.cos(thetas[p]), s=50, color='b')

    # Get the quadrature type
    quad_type = ''
    if quad.getQuadratureType() is openmoc.TABUCHI_YAMAMOTO:
        quad_type = 'TABUCHI_YAMAMOTO'
        title = 'TABUCHI YAMAMOTO'
    elif quad.getQuadratureType() is openmoc.LEONARD:
        quad_type = 'LEONARD'
        title = 'LEONARD'
    elif quad.getQuadratureType() is openmoc.GAUSS_LEGENDRE:
        quad_type = 'GAUSS_LEGENDRE'
        title = 'GAUSS LEGENDRE'
    elif quad.getQuadratureType() is openmoc.EQUAL_WEIGHT:
        quad_type = 'EQUAL_WEIGHT'
        title = 'EQUAL WEIGHT'
    elif quad.getQuadratureType() is openmoc.EQUAL_ANGLE:
        quad_type = 'EQUAL_ANGLE'
        title = 'EQUAL ANGLE'
    elif quad.getQuadratureType() is openmoc.CUSTOM:
        quad_type = 'CUSTOM'
        title = 'CUSTOM'
    else:
        py_printf('ERROR', 'Unable to plot the quadrature since the ' +
                  'quadrature type could not be recognized')

    title += ' with ' + str(num_azim) + ' azim, ' + \
             '{:5.3f}'.format(azim_spacing) + ' spacing and ' \
             + str(2*num_polar_2) + ' polar'

    filename = directory + 'quad_' + quad_type + '_' + \
               str(num_azim) + '_azim_' + '{:5.3f}'.format(azim_spacing) + \
               '_cm_spacing_' + str(2*num_polar_2) + '_polar.png'

    ax.view_init(elev=30, azim=45)
    ax.set_xlim([0,1])
    ax.set_ylim([0,1])
    ax.set_zlim([0,1])
    plt.title(title)

    # Restore settings if called from ipython
    matplotlib.rcParams.update(curr_rc)

    # Save the figure or return to user
    if track_generator.getGeometry().isRootDomain():
        if get_figure:
            return fig
        else:
            fig.savefig(filename, bbox_inches='tight')
            plt.close(fig)


class PlotParams(object):
    """A container for plotting parameters for spatially-varying plots.

    Attributes
    ----------
    geometry : openomc.Geometry
        The Geometry to query when generating the spatial map
    domain_type : str
        The domain type used to map spatial data to the geometry
    filename : str
        The filename string
    extension : str
        The image file extension (e.g., '.png')
    library : {'matplotlib', 'pil'}
        The image processing library to use to generate the plot
    offset : Real, optional
        The offset in the axis perpendicular to the plane (default is 0.0)
    gridsize : Integral
        The number of points along the x- and y-axes
    xlim : 2-tuple of Real
        A 2-tuple of (maximum, minimum) x-coordinates to display
    ylim : 2-tuple of Real
        A 2-tuple of (maximum, minimum) y-coordinates to display
    zlim : 2-tuple of Real, optional
        The minimim/maximum z-coordinates
    plane : {'xy', 'xz', 'yz'}
        The plane across which data is extracted and plotted
    offset : Real, optional
        The location of the plotted plane on the perpendicular axis
    title : str
        The minor title string
    suptitle : str
        The major title string
    norm : bool
        Normalize the plotted data to unity
    transparent_zeros : bool
        Make zeros in the data appear transparent
    interpolation : str
        Interpolation used between points (e.g., 'nearest')
    colorbar : bool
        Include a colorbar to the right of the plot
    cmap : matplotlib.colormap
        A Matplotlib colormap for the plot
    vmin : Real
        The minimum value used in colormapping the data
    vmax : Real
        The maximum value used in colormapping the data

    """

    def __init__(self):

        self._geometry = None
        self._domain_type = 'fsr'
        self._filename = None
        self._extension = '.png'
        self._library = 'matplotlib'
        self._plane = 'xy'
        self._offset = 0
        self._gridsize = 250
        self._xlim = None
        self._ylim = None
        self._zlim = None
        self._title = None
        self._suptitle = None
        self._norm = False
        self._transparent_zeros = False
        self._interpolation = None
        self._colorbar = False
        self._cmap = plt.get_cmap('nipy_spectral')
        self._vmin = None
        self._vmax = None

    @property
    def geometry(self):
        return self._geometry

    @property
    def domain_type(self):
        return self._domain_type

    @property
    def filename(self):
        return self._filename

    @property
    def extension(self):
        return self._extension

    @property
    def library(self):
        return self._library

    @property
    def plane(self):
        return self._plane

    @property
    def offset(self):
        return self._offset

    @property
    def gridsize(self):
        return self._gridsize

    @property
    def xlim(self):
        return self._xlim

    @property
    def ylim(self):
        return self._ylim

    @property
    def zlim(self):
        return self._zlim

    @property
    def colorbar(self):
        return self._colorbar

    @property
    def title(self):
        return self._title

    @property
    def suptitle(self):
        return self._suptitle

    @property
    def norm(self):
        return self._norm

    @property
    def transparent_zeros(self):
        return self._transparent_zeros

    @property
    def interpolation(self):
        return self._interpolation

    @property
    def cmap(self):
        return self._cmap

    @property
    def vmin(self):
        return self._vmin

    @property
    def vmax(self):
        return self._vmax

    @geometry.setter
    def geometry(self, geometry):
        cv.check_type('geometry', geometry, openmoc.Geometry)
        self._geometry = geometry
        self._check_offset()

    @domain_type.setter
    def domain_type(self, domain_type):
        cv.check_value('domain_type', domain_type, ('material', 'cell', 'fsr'))
        self._domain_type = domain_type

    @filename.setter
    def filename(self, filename):
        cv.check_type('filename', filename, basestring)
        self._filename = filename

    @extension.setter
    def extension(self, extension):
        cv.check_type('extension', extension, basestring)
        self._extension = extension

    @library.setter
    def library(self, library):
        cv.check_value('library', library, ('matplotlib', 'pil'))
        self._library = library

    @plane.setter
    def plane(self, plane):
      cv.check_value('plane', plane, ('xy', 'xz', 'yz'))
      self._plane = plane

    @offset.setter
    def offset(self, offset):
        if offset:
            self._offset = offset
        self._check_offset()

    @gridsize.setter
    def gridsize(self, gridsize):
        cv.check_type('gridsize', gridsize, Integral)
        cv.check_greater_than('gridsize', gridsize, 0)
        self._gridsize = gridsize

    @xlim.setter
    def xlim(self, xlim):
        if xlim:
            cv.check_type('xlim', xlim, tuple)
            cv.check_length('xlim', xlim, 2, 2)
        self._xlim = xlim

    @ylim.setter
    def ylim(self, ylim):
        if ylim:
            cv.check_type('ylim', ylim, tuple)
            cv.check_length('ylim', ylim, 2, 2)
        self._ylim = ylim

    @zlim.setter
    def zlim(self, zlim):
        if zlim:
            cv.check_type('zlim', zlim, tuple)
            cv.check_length('zlim', zlim, 2, 2)
            self._zlim = zlim

    @colorbar.setter
    def colorbar(self, colorbar):
        cv.check_type('colorbar', colorbar, bool)
        self._colorbar = colorbar

    @title.setter
    def title(self, title):
        cv.check_type('title', title, basestring)
        self._title = title

    @suptitle.setter
    def suptitle(self, suptitle):
        cv.check_type('suptitle', suptitle, basestring)
        self._suptitle = suptitle

    @norm.setter
    def norm(self, norm):
        cv.check_type('norm', norm, bool)
        self._norm = norm

    @transparent_zeros.setter
    def transparent_zeros(self, transparent_zeros):
        cv.check_type('transparent_zeros', transparent_zeros, bool)
        self._transparent_zeros = transparent_zeros

    @interpolation.setter
    def interpolation(self, interpolation):
        cv.check_type('interpolation', interpolation, basestring)
        self._interpolation = interpolation

    @cmap.setter
    def cmap(self, cmap):
        cv.check_type('cmap', cmap, (matplotlib.colors.Colormap, type(None)))
        self._cmap = cmap

    @vmin.setter
    def vmin(self, vmin):
        cv.check_type('vmin', vmin, Real)
        self._vmin = vmin

    @vmax.setter
    def vmax(self, vmax):
        cv.check_type('vmax', vmax, Real)
        self._vmax = vmax

    def _check_offset(self):
        if self.offset and self.geometry and self.plane:
            cv.check_type('offset', self.offset, Real)
            root = self.geometry.getRootUniverse()
            if self.plane == 'xy':
                cv.check_greater_than('offset', self.offset,
                                      root.getMinZ(), equality=True)
                cv.check_less_than('offset', self.offset,
                                   root.getMaxZ(), equality=True)
            if self.plane == 'xz':
                cv.check_greater_than('offset', self.offset,
                                      root.getMinY(), equality=True)
                cv.check_less_than('offset', self.offset,
                                  root.getMaxY(), equality=True)
            if self.plane == 'yz':
                cv.check_greater_than('offset', self.offset,
                                      root.getMinX(), equality=True)
                cv.check_less_than('offset', self.offset,
                                  root.getMaxX(), equality=True)


def _get_pixel_coords(plot_params):
    """A helper method to define coordinates for a plotting window.

    This routine builds a coordinate surface map for the plotting window defined
    for by the user. If no window was defined, then this routine uses the outer
    bounding box around the geometry as the plotting window.

    Parameters
    ----------
    plot_params : openmoc.plotter.PlotParams
        A PlotParams object initialized with a geometry

    Returns
    -------
    coords : dict
        A dictionary with the plotting window map and bounding box

    """

    # initialize variables to be returned
    geometry = plot_params.geometry
    coords = dict()
    bounds = dict()

    root = geometry.getRootUniverse()
    bounds['x'] = [root.getMinX() + TINY_MOVE,
                   root.getMaxX() - TINY_MOVE]
    bounds['y'] = [root.getMinY() + TINY_MOVE,
                   root.getMaxY() - TINY_MOVE]
    bounds['z'] = [root.getMinZ() + TINY_MOVE,
                   root.getMaxZ() - TINY_MOVE]

    if not plot_params.xlim is None:
        bounds['x'] = list(plot_params.xlim)

    if not plot_params.ylim is None:
        bounds['y'] = list(plot_params.ylim)

    if not plot_params.zlim is None:
        bounds['z'] = list(plot_params.zlim)


    # add attributes to coords dictionary
    if plot_params.plane == 'xy':
        coords['dim1'] = np.linspace(bounds['x'][0], bounds['x'][1],
                                     plot_params.gridsize)
        coords['dim2'] = np.linspace(bounds['y'][0], bounds['y'][1],
                                     plot_params.gridsize)
        coords['bounds'] = bounds['x'] + bounds['y']
    elif plot_params.plane == 'xz':
        coords['dim1'] = np.linspace(bounds['x'][0], bounds['x'][1],
                                     plot_params.gridsize)
        coords['dim2'] = np.linspace(bounds['z'][0], bounds['z'][1],
                                     plot_params.gridsize)
        coords['bounds'] = bounds['x'] + bounds['z']
    elif plot_params.plane == 'yz':
        coords['dim1'] = np.linspace(bounds['y'][0], bounds['y'][1],
                                     plot_params.gridsize)
        coords['dim2'] = np.linspace(bounds['z'][0], bounds['z'][1],
                                     plot_params.gridsize)
        coords['bounds'] = bounds['y'] + bounds['z']


    return coords


def _colorize(data, num_colors, seed=1):
    """Replace unique data values with a random but reproducible color IDs.

    This routine randomly assigns a random integer to each unique value in the
    input array. This is a helper routine for the plotting routine.

    Parameters
    ----------
    data : numpy.ndarray
        A NumPy array of data to colorize
    num_colors : Integral
        The number of random colors to generate
    seed : Integral
        The random number seed used to generate colors

    Returns
    -------
    ids_to_colors : numpy.ndarray
        A NumPy array with the random integer colors

    """

    # Generate linearly-spaced array of color indices
    all_ids = np.arange(num_colors, dtype=np.int64)

    # Generate linearly-spaced integer color IDs
    id_colors = np.arange(num_colors, dtype=np.int64)

    # Randomly shuffle the linearly-spaced integer color IDs
    numpy.random.seed(seed)
    np.random.shuffle(id_colors)

    # Insert random colors into appropriate locations in data array
    ids_to_colors = np.arange(num_colors, dtype=np.int64)
    ids_to_colors[all_ids] = id_colors

    return ids_to_colors.take(data)


def _get_pil_image(array, plot_params):
    """Plot 2D NumPy array data using Python Imaging Library (PIL).

    This is a good alternative to matplotlib for high-resolution images.

    Parameters
    ----------
    array : numpy.ndarray
        A NumPy array of data
    plot_params : openmoc.plotter.PlotParmas
        A PlotParams object with the matplotlib colormap to use

    Returns
    -------
    float_array : PIL.Image
        A Python Imaging Library (PIL) Image object

    """

    from PIL import Image

    # Convert array to a normalized array of floating point values
    float_array = np.zeros(array.shape, dtype=np.float)
    float_array[:,:] = array[:,:]
    float_array[:,:] /= np.max(float_array)

    # Use Python Imaging Library (PIL) to create an image from the array
    return Image.fromarray(np.uint8(plot_params.cmap(float_array) * 255))
